# coding=utf-8
# 获取miRbase 的数据 进行过滤的第一步 reads 的过滤
# 条件 mature star的长度在20~26nt之间
# 5'末端同质化

import numpy as np
from Bio import SeqIO
class pre(object):
    def __init__(self,name,seq):
        self.name = name
        self.seq = seq
        self.array = np.zeros(len(self.seq))
        self.T = True
        self.reason = ''
        self.count = 0
    def fetch_bwt(self, bowtieLista):
        count = 1
        try:
            count = int(bowtieLista[0].upper().split('_X')[1])
        except:
            print("reads name should be xxx_X123 like")
        length = len(bowtieLista[4])
        start =  int(bowtieLista[3])
        self.count += count
        self.array[start]+=count
        # make lenght to array
    def check_array(self):
        if self.count < 100:
            self.T = False
            self.reason = 'do not have enough reads matched'
        # 计算reads是否满足小RNA的特性
        # 先查５ｐ 
        max = np.max(self.array[0:int(len(self.array)/2)-5])
        inconsisit_5 = 0
        for i in range(0,int(len(self.array)/2)-5):
            if (max-self.array[i])<int(max/1.5):
                inconsisit_5+=1
        if inconsisit_5>2:
            self.reason = '5-end inhomogeneity of miR-5p'
            self.T = False
            return 0
        if max ==0:
            self.reason = 'no reads map to miR-5p'
            self.T = False
            return 0
        # 查３ｐ
        max = np.max(self.array[int(len(self.array)/2)-5:])
        inconsisit_3 = 0
        for i in range(int(len(self.array)/2)-5,int(len(self.array))):
            if (max-self.array[i])<int(max/1.5):
                inconsisit_3+=1
        if inconsisit_3>2:
            self.reason = '3-end inhomogeneity of miR-3p'
            self.T = False
            return 0
        if max ==0:
            self.reason = 'no reads map to miR-3p'
            self.T = False
            return 0


# 输入文件 miRbase上下载的前体序列
pre_file = r'pre.fa'
# bowtie的比对结果
bowtie_file = r'bwtOut'


dic = {}
for i in SeqIO.parse(pre_file,'fasta'):
    seq = str(i.seq)
    name = str(i.name)
    dic[name] = pre(name,seq)
# 读取bowtie文件
with open(bowtie_file) as fila:
    for i in fila:
        k = i.replace('\n','').split('\t')
        if len(k)>1 and k[1]=='+':
            name = k[2]
            dic[name].fetch_bwt(k)

outfile1 = 'discard_miRNA_from_miRbase'

handel1 = open(outfile1,'w')


for i in dic:
    dic[i].check_array()
    if dic[i].T != True:
        handel1.write(dic[i].name+'\t'+dic[i].reason+'\n')


handel1.close()


